# Author: qsj
# Time: 2023/02/24
# Describe: The code of Whale Optimization Algorithm(consult: https://seyedalimirjalili.com/woa)

import math
import random

import numpy
from matplotlib import pyplot as plt

from testFunctions import f1, f5


def WhaleOptimizationAlgorithm(n=100, d=2, T=100, ub=None, lb=None, f=None):
    """
    Whale Optimization Algorithm(WOA)

    n: Number of population
    d: Individual position dimension(Solution dimension)
    T: Maximum Iterations
    ub: Upper boundary list
    lb: Lower boundary list
    f: Fitness function

    1. Please modify the calculation part of fitness function (y)
    2. Default is minimum optimization, change "best_score" and "y < best_score" to maximum optimization
    3.The return value is an iterative curve, which can be modified to the optimal solution if necessary
    """

    best_index = 0
    best_score = 1e10  # Find min value use inf

    # a list use to record best position of each iteration
    Convergence_curve = []

    # position initialization
    X = [[random.uniform(lb[i], ub[i]) for i in range(d)] for j in range(n)]
    X = numpy.asarray(X)

    for t in range(T):
        for i, x in enumerate(X):

            # Check bounds
            for j in range(len(x)):
                if x[j] > ub[j]:
                    x[j] = ub[j]
                elif x[j] < lb[j]:
                    x[j] = lb[j]
                else:
                    continue

            # Calculate fitness
            y = f(x)

            # Updata best position
            if y < best_score:  # Find min value use <
                best_score = y
                best_index = i

        # param
        a = 2 * (1 - t / T)
        a2 = -1 - t / T

        #
        positions = X.copy()

        # Updata position of agent
        for i, x in enumerate(X):

            # random seed
            r1 = random.random()
            r2 = random.random()

            # Param of "Surrounding predation stage"
            A = 2 * a * r1 - a
            C = 2 * r2

            # Param of "Bubble attack stage"
            g = (a2 - 1) * random.random() + 1
            p = random.random()

            for j in range(d):
                if p < 0.5:  # Shrink Surround
                    if abs(A) >= 1:  # Random walk
                        rand_index = random.randint(0, n - 1)
                        D_rand = abs(C * X[rand_index][j] - x[j])
                        positions[i][j] = X[rand_index][j] - A * D_rand
                    else:  # Walk towards the best position
                        D = abs(C * X[best_index][j] - x[j])
                        positions[i][j] = X[best_index][j] - A * D
                else:  # Spiral update position
                    b = 1
                    Dl = abs(X[best_index][j] - x[j])
                    positions[i][j] = Dl * math.exp(b * g) * math.cos(2 * math.pi * g) + X[best_index][j].copy()

        # Replace
        X = positions.copy()

        # Record best position to updata iterative process curve
        Convergence_curve.append(best_score)

    # Return optimal solution
    return Convergence_curve


# # example
# population_numbers = 100
# dimension = 20
# iterations = 100
#
# upper_bounds = [30 for i in range(dimension)]
# lower_bounds = [-30 for i in range(dimension)]
#
# xl = WhaleOptimizationAlgorithm(n=population_numbers, d=dimension, T=iterations,
#                                 ub=upper_bounds, lb=lower_bounds, f=f5)
#
# # Display the optimal result
# print("Optimization result: ", xl[-1])
#
# # Draw iterative change process curve
# plt.title('WOA Iterative Process')
# plt.xlabel('iterations')
# plt.ylabel('best score')
# plt.plot([i for i in range(iterations)], xl)
# plt.show()
